import numpy as np
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from openpyxl import Workbook


def fisher(X, y):
    """
    使用Fisher判别法进行线性判别分析

    参数：
    X: 输入特征矩阵，每行代表一个样本，每列代表一个特征
    y: 标签向量，每个元素代表对应样本的类别

    返回值：
    w: 判别函数系数矩阵，每行代表一个判别函数的系数向量
    """
    n_sum = y.shape[0]
    classes = np.unique(y)
    n_features = X.shape[1]
    n_classes = len(classes)
    # 计算每个类别的均值向量
    class_means = np.array([np.mean(X[y == c], axis=0) for c in classes])

    # 计算全局均值向量
    overall_mean = np.mean(X, axis=0)

    # 计算类内散布矩阵
    Sw = np.zeros((n_features, n_features))
    for c in classes:
        class_samples = X[y == c]
        diff = class_samples - class_means[c]
        Sw += np.dot(diff.T, diff)

    # 计算类间散布矩阵
    Sb = np.zeros((n_features, n_features))
    for c in classes:
        ni = len(X[y == c])
        diff = class_means[c] - overall_mean
        Sb += ni * np.outer(diff, diff)

    # 计算Fisher准则函数的系数矩阵
    eigvals, eigvecs = np.linalg.eig(np.linalg.inv(Sw) @ Sb)
    # 选取特征值非零的特征向量并按照特征值大小排序
    non_zero_indices = np.where(eigvals > 1e-10)[0]
    sorted_indices = np.argsort(eigvals[non_zero_indices])[::-1]
    w = eigvecs[:, non_zero_indices[sorted_indices]].real
    # 标准化判别系数
    w_normalized = w
    for i in range(w_normalized.shape[1]):
        w_normalized[:, i] = w_normalized[:, i] / np.sqrt(
            w_normalized[:, i].T @ (1 / (n_sum - n_classes) * Sw) @ w_normalized[:, i]
        )
    print("计算完毕，判别系数为")
    print(w_normalized)
    print("贡献率分别为")
    print(eigvals[sorted_indices].real / np.sum(eigvals[sorted_indices].real))
    return w_normalized


def write_array_to_excel(array, filename):
    # 创建一个新的工作簿
    wb = Workbook()
    # 激活第一个工作表
    ws = wb.active

    # 遍历二维数组，并将其写入工作表中
    for row in array:
        ws.append(row)

    # 保存工作簿到指定的文件
    wb.save(filename)
    print(f"二维数组已写入到 {filename} 文件中.")


if __name__ == "__main__":
    # 示例用法
    from data import X, y

    # 调用费希尔判别法函数
    w = fisher(X, y)

    # 使用判别函数计算
    print("使用判别函数计算的结果：")
    print(X @ w)
    write_array_to_excel((X @ w).tolist(), "fisher.xlsx")
    # 调库验证
    # 创建并训练Fisher判别器
    # lda = LinearDiscriminantAnalysis()
    # lda.fit(X, y)

    # # 输出判别函数系数
    # print("判别函数系数矩阵：")
    # print(lda.coef_)
